import pandas as pd
from sklearn.preprocessing import LabelEncoder
import matplotlib.pyplot as plt
import seaborn as sns

# 加载训练数据
train_data = pd.read_csv('../data/train.csv', header=None)

# 自动将所有可以转换为数字的列转为数值类型
train_data = train_data.apply(pd.to_numeric, errors='coerce')


# 添加列名
columns = ['Attrition', 'Age', 'BusinessTravel', 'Department', 'DistanceFromHome', 'Education', 'EducationField',
           'EmployeeNumber', 'EnvironmentSatisfaction', 'Gender', 'JobInvolvement', 'JobLevel', 'JobRole',
           'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome', 'NumCompaniesWorked', 'Over18', 'OverTime',
           'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction', 'StandardHours', 'StockOptionLevel',
           'TotalWorkingYears', 'TrainingTimesLastYear', 'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole',
           'YearsSinceLastPromotion', 'YearsWithCurrManager']
train_data.columns = columns

# 编码类别型变量
categorical_features = ['BusinessTravel', 'Department', 'EducationField', 'Gender', 'JobRole', 'MaritalStatus', 'Over18', 'OverTime']
label_encoders = {}
for feature in categorical_features:
    le = LabelEncoder()
    train_data[feature] = le.fit_transform(train_data[feature])
    label_encoders[feature] = le

# 将 Attrition 转换为二进制标签
train_data['Attrition'] = train_data['Attrition'].apply(lambda x: 1 if x == 1 else 0)

# 年龄分布
plt.figure(figsize=(10, 6))
sns.histplot(data=train_data, x='Age', hue='Attrition', bins=20, kde=True)
plt.title('Age Distribution by Attrition')
plt.xlabel('Age')
plt.ylabel('Count')
plt.show()

# 薪资分布
plt.figure(figsize=(10, 6))
sns.boxplot(x='Attrition', y='MonthlyIncome', data=train_data)
plt.title('Monthly Income Distribution by Attrition')
plt.xlabel('Attrition')
plt.ylabel('Monthly Income')
plt.show()

# 工作年限分布
plt.figure(figsize=(10, 6))
sns.boxplot(x='Attrition', y='TotalWorkingYears', data=train_data)
plt.title('Total Working Years Distribution by Attrition')
plt.xlabel('Attrition')
plt.ylabel('Total Working Years')
plt.show()

# 满意度分布
plt.figure(figsize=(10, 6))
sns.boxplot(x='Attrition', y='JobSatisfaction', data=train_data)
plt.title('Job Satisfaction Distribution by Attrition')
plt.xlabel('Attrition')
plt.ylabel('Job Satisfaction')
plt.show()

# 相关性热力图
corr_matrix = train_data.corr()
plt.figure(figsize=(24, 16))
sns.heatmap(corr_matrix, annot=True, cmap='coolwarm')
plt.title('Correlation Matrix')
plt.show()